package com.haibara.codesandbox.pool.containerPool;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.exception.NotFoundException;
import com.github.dockerjava.api.model.Bind;
import com.github.dockerjava.api.model.Capability;
import com.github.dockerjava.api.model.HostConfig;
import com.haibara.codesandbox.constant.FileConstant;
import jakarta.annotation.PreDestroy;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.actuate.autoconfigure.wavefront.WavefrontProperties;
import org.springframework.stereotype.Component;

import java.util.ArrayDeque;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 容器池
 *
 * @Author: mengyu
 * @Date: 2025/6/22
 */
@Component
@Slf4j
public class ContainerPool {
    private final ConcurrentHashMap<Long, ArrayDeque<CreateContainerResponse>> containerMap = new ConcurrentHashMap<>();
    private final int TOTAL_CONTAINER_SIZE = 150;
    private AtomicInteger totalCount = new AtomicInteger(0);
    private final ReentrantLock lock = new ReentrantLock();

    @Resource
    private DockerClient dockerClient;


    public CreateContainerResponse getContainer(Long questionId) {
        lock.lock();
        try {
            containerMap.putIfAbsent(questionId, new ArrayDeque<>());
            ArrayDeque<CreateContainerResponse> containerDeque = containerMap.get(questionId);
            if (containerDeque.isEmpty()) {
                // 容器队列为空
                CreateContainerResponse container = dockerClient.createContainerCmd("alpine:3.12.0")
                        .withTty(true)
                        .withAttachStdin(true)
                        .withAttachStderr(true)
                        .withAttachStdout(true)
                        // 配置数据卷
                        .withHostConfig(getHostConfig(questionId))
                        .exec();
                dockerClient.startContainerCmd(container.getId()).exec();
                return container;
            } else {
                // 容器队列不为空
                totalCount.decrementAndGet();
                return containerDeque.pollFirst();
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 返还容器
     *
     * @param container  待返还的容器
     * @param questionId 题目id
     */
    public void returnContainer(CreateContainerResponse container, Long questionId) {
        lock.lock();
        try {
            if (totalCount.get() <= TOTAL_CONTAINER_SIZE) {
                // 容器池未满，可以返还
                containerMap.putIfAbsent(questionId, new ArrayDeque<>());
                ArrayDeque<CreateContainerResponse> containerDeque = containerMap.get(questionId);
                containerDeque.push(container);
                totalCount.incrementAndGet();
            } else {
                deleteContainer(container);
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 清除所有容器
     */
    public void clearAllContainer() {
        // 清理容器
        for (Long questionId : containerMap.keySet()) {
            ArrayDeque<CreateContainerResponse> containerDeque = containerMap.get(questionId);
            log.info("题目id：{}的容器队列长度为：{}", questionId, containerDeque.size());
            for (CreateContainerResponse createContainerResponse : containerMap.get(questionId)) {
                deleteContainer(createContainerResponse);
            }
            containerDeque.clear();
        }
        totalCount = new AtomicInteger(0);
    }

    public String getPoolInfo() {
        lock.lock();
        try {
            StringBuilder sb = new StringBuilder("容器池状态：\n");
            int totalContainers = 0;
            for (Long questionId : containerMap.keySet()) {
                int queueSize = containerMap.get(questionId).size();
                totalContainers += queueSize;
                sb.append(String.format("题目ID %d: 可用容器%d个\n", questionId, queueSize));
            }
            if (totalContainers > 0) {
                int colonIndex = sb.indexOf("：");
                sb.insert(colonIndex + 2, "总容器数=" + totalContainers + ", ");
            } else {
                sb.append("容器池暂无容器");
            }
            sb.append(String.format("容器池剩余容量：%d个\n", TOTAL_CONTAINER_SIZE - totalCount.get()));
            log.info("{}", sb);
            return sb.toString();
        } finally {
            lock.unlock();
        }
    }

    private HostConfig getHostConfig(Long questionId) {
        // 2.1.1 配置数据卷
        HostConfig hostConfig = HostConfig.newHostConfig();
        hostConfig.withBinds(
                // 用户代码数据卷
                Bind.parse(FileConstant.codeFilePath + ":" + FileConstant.codeFilePath),
                // 测试文件数据卷，设定为只读
                Bind.parse(FileConstant.inputFilePath + "/" + questionId + ":" + FileConstant.inputFilePath + "/" + questionId + ":ro"),
                // 文件输出数据卷
                Bind.parse(FileConstant.userOutputFilePath + ":" + FileConstant.userOutputFilePath)
        );
        // 2.1.2 配置容器限制
        hostConfig
                .withCpuShares(512)
                // 禁止特权命令
                .withPrivileged(false)
                .withCapDrop(Capability.valueOf("ALL"))
                // 设置容器内部文件为只读
                .withReadonlyRootfs(true)
                // 禁止容器联网
                .withNetworkMode("none")
                // 设置最大内存
                .withMemory(262144000L);
        return hostConfig;
    }

    private void deleteContainer(CreateContainerResponse container) {
        // 删除容器
        log.info("同步删除容器开始{}", container.getId());
        try {
            dockerClient.removeContainerCmd(container.getId())
                    .withForce(true)
                    .exec();
            log.info("同步删除容器结束{}", container.getId());
        } catch (NotFoundException e) {
            log.info("同步删除容器失败{}", e.getMessage());
        }
    }

    @PreDestroy
    public void cleanContainerBeforeDestroy() {
        log.info("开始执行容器池销毁前清理");
        clearAllContainer();
        log.info("容器池清理完成");
    }
}